package com.project.srtp.common.config;

import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.concurrent.ConcurrentHashMap;

/**
 * @author 曹行阳
 * @version 0.1
 * @date 2025/6/17 1:38
 */

@Component
public class WebSocketHandler extends TextWebSocketHandler {

    // userId -> session
    private static final ConcurrentHashMap<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String userId = getUserId(session);
        if (userId != null) {
            userSessions.put(userId, session);
            System.out.println("用户上线：" + userId);
        } else {
            session.close();
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // JSON payload: { "senderId": "u1", "receiverId": "u2", "cipher": "...", ... }
        String payload = message.getPayload();
        WebSocketSession receiverSession = userSessions.get(getReceiverId(payload));

        if (receiverSession != null && receiverSession.isOpen()) {
            receiverSession.sendMessage(new TextMessage(payload));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String userId = getUserId(session);
        if (userId != null) {
            userSessions.remove(userId);
            System.out.println("用户下线：" + userId);
        }
    }

    private String getUserId(WebSocketSession session) {
        // ws://localhost:8080/ws?userId=123
        String query = session.getUri().getQuery(); // userId=123
        if (query != null && query.startsWith("userId=")) {
            return query.substring(7);
        }
        return null;
    }

    private String getReceiverId(String json) {
        try {
            com.fasterxml.jackson.databind.ObjectMapper mapper = new com.fasterxml.jackson.databind.ObjectMapper();
            return mapper.readTree(json).get("receiverId").asText();
        } catch (Exception e) {
            return null;
        }
    }
}

